import os
import argparse
import yaml
from utils import *
from gatys_nst import *


def parse():
    parser = argparse.ArgumentParser(description="Neural style transfer (Leon A. Gatys' method)")
    parser.add_argument("--config", default="config.yml", type=str, required=False)
    return parser.parse_args()


def main():
    opt = parse()
    with open(opt.config, encoding="utf-8") as fp:
        config = yaml.load(fp, Loader=yaml.FullLoader)
        config = config["params"]

    # The basic parameters
    device = config["device"]
    content_layers, style_layers = config["content_layers"], config["style_layers"]
    content_image, style_image = Image.open(config["content_image"]), Image.open(config["style_image"])

    # Initialize data for the pretraining network
    image_shape = content_image.size if config["image_shape"] is None else config["image_shape"]
    network = get_pretrained_net(config["network_name"], content_layers, style_layers, device)
    content_x, content_y = get_content_data(content_image, image_shape, network, content_layers, device)
    style_x, style_gram_y = get_style_data(style_image, image_shape, network, style_layers, device)

    # Initialize to generate an image
    begin_epoch = 0  # Training starts from 0 by default
    init_image = config["init_image"].lower()
    init_image = "content" if init_image is None else init_image
    init_image_values = ("content", "style", "continue", "random", "combine", "default")
    assert init_image in init_image_values, f"The init_image option only can be {', '.join(init_image_values)}."
    if init_image == "content" or init_image == "default":
        init_image = content_x
    elif init_image == "style":
        init_image = style_x
    elif init_image == "continue":
        checkpoints = f"./checkpoints/{config['work_name']}"
        assert config["work_name"] is not None and os.path.exists(checkpoints), \
            "When init_image is continue, the value of the work_name item cannot be empty."
        last_image_box = sorted(os.listdir(checkpoints), reverse=True)
        if last_image_box is None or len(last_image_box) == 0:
            init_image = content_x
        else:
            last_image = Image.open(f"{checkpoints}/{last_image_box[0]}")
            init_image, _ = get_content_data(last_image, image_shape, network, content_layers, device)
            begin_epoch = int(last_image_box[0].split(".")[0]) + 1
    elif init_image == "combine":
        assert config["combine_style_weight"] is not None and config["combine_content_weight"] is not None, \
            "When the combine option is specified, the 'combine_style_weight' and 'combine_content_weight' parameters cannot be empty."
        init_image = content_x * config["combine_content_weight"] + style_x * config["combine_style_weight"]
    else:
        init_image = None

    train(content_y=content_y, style_gram_y=style_gram_y, num_epochs=config["num_epochs"], lr=config["lr"],
          lr_decay_epoch=config["lr_decay_epoch"], channels=3, image_nrows=image_shape[0], image_ncols=image_shape[1],
          network=network, content_layers=content_layers, style_layers=style_layers,
          content_weight=config["content_weight"], style_weight=config["style_weight"], tv_weight=config["tv_weight"],
          begin_epoch=begin_epoch, init_image=init_image, device=device, step=config["step"],
          work_name=config["work_name"])


if __name__ == '__main__':
    main()
